"""
BiasCorrector Core Types and Algorithms

This module implements the fundamental bias correction infrastructure for atmospheric 
data assimilation, providing core data structures, configuration management, and 
basic mathematical operations for systematic bias removal in observations.

Based on GSI's bias correction system with support for:
- Variational bias correction (VarBC) for satellite radiances
- Air mass bias correction for conventional observations  
- Angle-dependent bias correction for satellite wind data
- Static and adaptive bias correction schemes
"""
module BiasCorrector

using LinearAlgebra
using SparseArrays
using Statistics
using Dates
using ..ObservationTypes

export BiasScheme, BiasConfiguration, BiasState, BiasStatistics
export VarBCScheme, AirMassScheme, AngleScheme, StaticScheme
export BiasPredictor, PredictorType, LinearPredictor, NonlinearPredictor
export BiasCorrection, ChannelBias, InstrumentBias, ObservationBias
export initialize_bias_correction, update_bias_state, apply_bias_correction
export compute_bias_predictors, validate_bias_configuration

"""
Abstract base type for bias correction schemes
"""
abstract type BiasScheme end

"""
Variational Bias Correction scheme for satellite radiances
Implements adaptive bias correction through variational minimization
"""
struct VarBCScheme <: BiasScheme
    channels::Vector{Int}
    predictors::Vector{Symbol}
    adaptation_window::Int
    convergence_threshold::Float64
    regularization_weight::Float64
    
    function VarBCScheme(channels, predictors; 
                        adaptation_window=50, 
                        convergence_threshold=1e-6,
                        regularization_weight=0.01)
        new(channels, predictors, adaptation_window, 
            convergence_threshold, regularization_weight)
    end
end

"""
Air mass bias correction for conventional observations
Corrects systematic biases related to atmospheric conditions
"""
struct AirMassScheme <: BiasScheme
    observation_types::Vector{Symbol}
    pressure_levels::Vector{Float64}
    temperature_ranges::Vector{Tuple{Float64, Float64}}
    humidity_ranges::Vector{Tuple{Float64, Float64}}
    correction_factors::Dict{Symbol, Matrix{Float64}}
end

"""
Angle-dependent bias correction for satellite wind data
Accounts for viewing angle effects in satellite observations
"""
struct AngleScheme <: BiasScheme
    satellite_ids::Vector{String}
    zenith_angle_bins::Vector{Float64}
    azimuth_angle_bins::Vector{Float64}
    correction_table::Dict{String, Matrix{Float64}}
    interpolation_method::Symbol # :linear, :cubic, :nearest
end

"""
Static bias correction using pre-computed tables
Simple lookup-based correction for well-characterized biases
"""
struct StaticScheme <: BiasScheme
    correction_table::Dict{Symbol, Dict{Int, Float64}}
    validity_period::Tuple{DateTime, DateTime}
    update_frequency::Period
end

"""
Enumeration of bias predictor types
"""
@enum PredictorType begin
    CONSTANT_PRED = 1
    LAPSE_RATE_PRED = 2
    LAPSE_RATE_SQUARED_PRED = 3
    COSINE_LATITUDE_PRED = 4
    SINE_LATITUDE_PRED = 5
    EMISSIVITY_PRED = 6
    SCAN_ANGLE_PRED = 7
    SCAN_ANGLE_SQUARED_PRED = 8
    CLOUD_LIQUID_WATER_PRED = 9
    CLOUD_ICE_WATER_PRED = 10
    SURFACE_TEMPERATURE_PRED = 11
    THICKNESS_PRED = 12
end

"""
Abstract base type for bias predictors
"""
abstract type BiasPredictor{T<:AbstractFloat} end

"""
Linear bias predictor implementation
Computes linear combinations of atmospheric/geometric variables
"""
struct LinearPredictor{T} <: BiasPredictor{T}
    predictor_type::PredictorType
    coefficients::Vector{T}
    scaling_factor::T
    offset::T
    valid_range::Tuple{T, T}
    
    function LinearPredictor{T}(ptype::PredictorType, coeffs::Vector{T};
                              scaling=one(T), offset=zero(T),
                              valid_range=(typemin(T), typemax(T))) where T
        new{T}(ptype, coeffs, scaling, offset, valid_range)
    end
end

"""
Nonlinear bias predictor with polynomial or exponential terms
Handles complex nonlinear bias relationships
"""
struct NonlinearPredictor{T} <: BiasPredictor{T}
    predictor_type::PredictorType
    polynomial_order::Int
    exponential_terms::Vector{T}
    interaction_terms::Matrix{T}
    normalization_params::Dict{Symbol, T}
    
    function NonlinearPredictor{T}(ptype::PredictorType, order::Int;
                                 exp_terms=T[], interactions=zeros(T, 0, 0),
                                 norm_params=Dict{Symbol, T}()) where T
        new{T}(ptype, order, exp_terms, interactions, norm_params)
    end
end

"""
Configuration for bias correction system
Contains all parameters and settings for bias correction schemes
"""
struct BiasConfiguration{T<:AbstractFloat}
    scheme_type::Type{<:BiasScheme}
    active_schemes::Vector{BiasScheme}
    predictor_config::Dict{Symbol, Vector{BiasPredictor{T}}}
    quality_control::Dict{Symbol, Any}
    output_diagnostics::Bool
    parallel_processing::Bool
    memory_limit::Int # MB
    
    function BiasConfiguration{T}(scheme_type::Type{<:BiasScheme};
                                schemes=BiasScheme[],
                                predictors=Dict{Symbol, Vector{BiasPredictor{T}}}(),
                                qc_params=Dict{Symbol, Any}(),
                                diagnostics=true,
                                parallel=true,
                                mem_limit=1024) where T
        new{T}(scheme_type, schemes, predictors, qc_params, 
               diagnostics, parallel, mem_limit)
    end
end

"""
Current state of bias correction system
Tracks bias coefficients, statistics, and convergence status
"""
mutable struct BiasState{T<:AbstractFloat}
    coefficients::Dict{Symbol, Matrix{T}}
    covariance_matrix::Dict{Symbol, Matrix{T}}
    iteration_count::Int
    convergence_history::Vector{T}
    last_update::DateTime
    is_converged::Bool
    
    function BiasState{T}() where T
        new{T}(Dict{Symbol, Matrix{T}}(),
               Dict{Symbol, Matrix{T}}(),
               0, T[], now(), false)
    end
end

"""
Bias correction statistics and quality metrics
Monitors correction performance and stability
"""
struct BiasStatistics{T<:AbstractFloat}
    mean_bias::Dict{Symbol, Vector{T}}
    bias_variance::Dict{Symbol, Vector{T}}
    correction_magnitude::Dict{Symbol, Vector{T}}
    temporal_stability::Dict{Symbol, T}
    spatial_coherence::Dict{Symbol, T}
    channel_correlations::Dict{Symbol, Matrix{T}}
    observation_count::Dict{Symbol, Int}
    last_computed::DateTime
end

"""
Channel-specific bias correction data
Stores bias information for individual satellite channels
"""
struct ChannelBias{T<:AbstractFloat}
    channel_number::Int
    instrument_id::String
    bias_coefficients::Vector{T}
    predictor_values::Vector{T}
    correction_value::T
    uncertainty::T
    quality_flag::Int
end

"""
Instrument-level bias correction container
Manages bias data for entire instruments or sensors
"""
struct InstrumentBias{T<:AbstractFloat}
    instrument_name::String
    satellite_id::String
    channel_biases::Dict{Int, ChannelBias{T}}
    global_offset::T
    calibration_date::DateTime
    validity_window::Tuple{DateTime, DateTime}
end

"""
Complete observation bias correction structure
Contains all bias information for a set of observations
"""
struct ObservationBias{T<:AbstractFloat}
    observation_type::ObservationType
    location::GeographicLocation{T}
    time_stamp::DateTime
    original_value::T
    corrected_value::T
    bias_correction::T
    correction_uncertainty::T
    applied_schemes::Vector{Symbol}
end

"""
Initialize bias correction system with configuration parameters

# Arguments
- `config`: BiasConfiguration containing all setup parameters
- `observation_types`: Vector of observation types to process
- `initial_coefficients`: Optional initial bias coefficients

# Returns
- Initialized BiasState ready for processing
"""
function initialize_bias_correction(config::BiasConfiguration{T},
                                   observation_types::Vector{ObservationType};
                                   initial_coefficients=nothing) where T
    state = BiasState{T}()
    
    # Initialize coefficient arrays for each observation type
    for obs_type in observation_types
        obs_symbol = Symbol(obs_type.name)
        
        if haskey(config.predictor_config, obs_symbol)
            predictors = config.predictor_config[obs_symbol]
            n_predictors = length(predictors)
            
            if obs_type isa RadianceObservation
                n_channels = length(obs_type.channels)
                if initial_coefficients !== nothing && haskey(initial_coefficients, obs_symbol)
                    state.coefficients[obs_symbol] = copy(initial_coefficients[obs_symbol])
                else
                    state.coefficients[obs_symbol] = zeros(T, n_channels, n_predictors)
                end
                state.covariance_matrix[obs_symbol] = Matrix{T}(I, n_channels * n_predictors, 
                                                               n_channels * n_predictors)
            else
                if initial_coefficients !== nothing && haskey(initial_coefficients, obs_symbol)
                    state.coefficients[obs_symbol] = copy(initial_coefficients[obs_symbol])
                else
                    state.coefficients[obs_symbol] = zeros(T, 1, n_predictors)
                end
                state.covariance_matrix[obs_symbol] = Matrix{T}(I, n_predictors, n_predictors)
            end
        end
    end
    
    state.last_update = now()
    return state
end

"""
Update bias correction state with new observations and analysis

# Arguments
- `state`: Current BiasState to update
- `observations`: Vector of processed observations
- `analysis_increment`: Analysis increment from data assimilation
- `config`: BiasConfiguration with update parameters

# Returns
- Updated BiasState with new coefficients and statistics
"""
function update_bias_state(state::BiasState{T}, 
                          observations::Vector{<:Observation{T}},
                          analysis_increment::Vector{T},
                          config::BiasConfiguration{T}) where T
    
    # Group observations by type for efficient processing
    obs_groups = group_observations_by_type(observations)
    
    convergence_metrics = T[]
    
    for (obs_type, obs_list) in obs_groups
        obs_symbol = Symbol(obs_type)
        
        if !haskey(state.coefficients, obs_symbol)
            continue
        end
        
        # Compute predictor matrix for this observation type
        predictor_matrix = compute_predictor_matrix(obs_list, 
                                                   config.predictor_config[obs_symbol])
        
        # Extract innovation vector (O - B) for bias correction
        innovation_vector = [obs.innovation for obs in obs_list]
        
        # Update bias coefficients using variational approach
        if any(isa(scheme, VarBCScheme) for scheme in config.active_schemes)
            old_coeffs = copy(state.coefficients[obs_symbol])
            state.coefficients[obs_symbol] = update_varbc_coefficients(
                state.coefficients[obs_symbol],
                predictor_matrix,
                innovation_vector,
                state.covariance_matrix[obs_symbol],
                config
            )
            
            # Compute convergence metric
            coeff_change = norm(state.coefficients[obs_symbol] - old_coeffs) / 
                          (norm(old_coeffs) + eps(T))
            push!(convergence_metrics, coeff_change)
        end
    end
    
    # Update state metadata
    state.iteration_count += 1
    state.last_update = now()
    
    # Check convergence
    if !isempty(convergence_metrics)
        mean_change = mean(convergence_metrics)
        push!(state.convergence_history, mean_change)
        
        # Convergence check with multiple criteria
        converged_schemes = filter(s -> isa(s, VarBCScheme), config.active_schemes)
        if !isempty(converged_schemes)
            threshold = converged_schemes[1].convergence_threshold
            state.is_converged = (mean_change < threshold && 
                                length(state.convergence_history) >= 3 &&
                                all(h < threshold * 2 for h in 
                                    state.convergence_history[end-2:end]))
        end
    end
    
    return state
end

"""
Apply bias correction to observations using current state

# Arguments
- `observations`: Vector of observations to correct
- `state`: Current BiasState with correction coefficients
- `config`: BiasConfiguration with correction parameters

# Returns
- Vector of ObservationBias with corrected values
"""
function apply_bias_correction(observations::Vector{<:Observation{T}},
                              state::BiasState{T},
                              config::BiasConfiguration{T}) where T
    corrected_obs = ObservationBias{T}[]
    
    for obs in observations
        obs_type = typeof(obs).parameters[1] # Get observation type
        obs_symbol = Symbol(obs_type)
        
        if !haskey(state.coefficients, obs_symbol)
            # No bias correction available, return uncorrected
            bias_obs = ObservationBias{T}(
                obs_type, obs.location, obs.time,
                obs.value, obs.value, zero(T), zero(T), Symbol[]
            )
            push!(corrected_obs, bias_obs)
            continue
        end
        
        # Compute predictor values for this observation
        predictors = compute_observation_predictors(obs, config.predictor_config[obs_symbol])
        
        # Apply bias correction
        bias_correction = zero(T)
        applied_schemes = Symbol[]
        
        if obs isa RadianceObservation
            channel_idx = findfirst(==(obs.channel), obs.channels)
            if channel_idx !== nothing
                bias_correction = dot(state.coefficients[obs_symbol][channel_idx, :], predictors)
                push!(applied_schemes, :varbc)
            end
        else
            bias_correction = dot(state.coefficients[obs_symbol][1, :], predictors)
            push!(applied_schemes, :conventional)
        end
        
        corrected_value = obs.value - bias_correction
        correction_uncertainty = estimate_correction_uncertainty(obs, state, config)
        
        bias_obs = ObservationBias{T}(
            obs_type, obs.location, obs.time,
            obs.value, corrected_value, bias_correction, 
            correction_uncertainty, applied_schemes
        )
        push!(corrected_obs, bias_obs)
    end
    
    return corrected_obs
end

"""
Compute bias predictors for a set of observations
"""
function compute_bias_predictors(observations::Vector{<:Observation{T}},
                               predictors::Vector{BiasPredictor{T}}) where T
    n_obs = length(observations)
    n_pred = length(predictors)
    predictor_matrix = zeros(T, n_obs, n_pred)
    
    for (i, obs) in enumerate(observations)
        for (j, predictor) in enumerate(predictors)
            predictor_matrix[i, j] = evaluate_predictor(predictor, obs)
        end
    end
    
    return predictor_matrix
end

"""
Validate bias correction configuration for consistency and completeness
"""
function validate_bias_configuration(config::BiasConfiguration{T}) where T
    validation_errors = String[]
    
    # Check scheme consistency
    for scheme in config.active_schemes
        if isa(scheme, VarBCScheme)
            if scheme.adaptation_window <= 0
                push!(validation_errors, "VarBC adaptation window must be positive")
            end
            if scheme.convergence_threshold <= 0
                push!(validation_errors, "VarBC convergence threshold must be positive")
            end
        end
    end
    
    # Validate predictor configuration
    for (obs_type, predictors) in config.predictor_config
        if isempty(predictors)
            push!(validation_errors, "No predictors defined for observation type $obs_type")
        end
        
        for predictor in predictors
            if isa(predictor, LinearPredictor) && isempty(predictor.coefficients)
                push!(validation_errors, "Linear predictor has empty coefficients for $obs_type")
            end
        end
    end
    
    # Check memory limits
    if config.memory_limit <= 0
        push!(validation_errors, "Memory limit must be positive")
    end
    
    return validation_errors
end

# Helper functions
function group_observations_by_type(observations)
    groups = Dict()
    for obs in observations
        obs_type = string(typeof(obs).parameters[1])
        if !haskey(groups, obs_type)
            groups[obs_type] = typeof(obs)[]
        end
        push!(groups[obs_type], obs)
    end
    return groups
end

function compute_predictor_matrix(observations, predictors)
    n_obs = length(observations)
    n_pred = length(predictors)
    matrix = zeros(eltype(observations[1].value), n_obs, n_pred)
    
    for (i, obs) in enumerate(observations)
        for (j, pred) in enumerate(predictors)
            matrix[i, j] = evaluate_predictor(pred, obs)
        end
    end
    
    return matrix
end

function compute_observation_predictors(observation, predictors)
    return [evaluate_predictor(pred, observation) for pred in predictors]
end

function evaluate_predictor(predictor::LinearPredictor{T}, observation) where T
    # Implementation depends on predictor type and observation content
    # This is a simplified version - full implementation would extract
    # specific atmospheric/geometric variables based on predictor_type
    return one(T) # Placeholder
end

function evaluate_predictor(predictor::NonlinearPredictor{T}, observation) where T
    # Nonlinear predictor evaluation with polynomial/exponential terms
    return one(T) # Placeholder
end

function update_varbc_coefficients(current_coeffs, predictor_matrix, innovations, 
                                 covariance, config)
    # Simplified VarBC update - full implementation would use
    # variational minimization with proper regularization
    return current_coeffs # Placeholder
end

function estimate_correction_uncertainty(observation, state, config)
    # Estimate uncertainty in bias correction based on coefficient covariance
    return 0.1 # Placeholder - should use actual uncertainty propagation
end

end # module BiasCorrector